package wsframework.wssecurity.utils;

import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import java.util.Vector;

import javax.xml.namespace.QName;

import org.apache.axiom.om.OMAttribute;
import org.apache.axiom.om.OMElement;
import org.apache.axiom.om.OMFactory;
import org.apache.axiom.om.OMNamespace;
import org.apache.axiom.om.xpath.AXIOMXPath;
import org.apache.axiom.soap.SOAPEnvelope;
import org.apache.axiom.soap.SOAPHeader;
import org.apache.axiom.soap.SOAPHeaderBlock;
import org.apache.axis2.saaj.util.SAAJUtil;
import org.apache.ws.security.WSConstants;
import org.apache.ws.security.WSEncryptionPart;
import org.apache.ws.security.WSSecurityEngineResult;
import org.apache.ws.security.WSUsernameTokenPrincipal;
import org.apache.ws.security.handler.WSHandlerResult;
import org.apache.ws.security.message.WSSecBase;
import org.apache.ws.security.message.WSSecEncryptedKey;
import org.apache.ws.security.util.WSSecurityUtil;
import org.apache.xml.security.utils.Constants;
import org.jaxen.JaxenException;
import org.jaxen.XPath;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;

import cartagows.wsframework.wssecurity.core.MessageData;
import cartagows.wsframework.wssecurity.core.SecurityException;
import cartagows.wsframework.wssecurity.policy.PolicyData;
import cartagows.wsframework.wssecurity.policy.SPConstants;
import cartagows.wsframework.wssecurity.policy.model.SupportingToken;
import cartagows.wsframework.wssecurity.policy.model.Wss10;
import cartagows.wsframework.wssecurity.policy.model.Wss11;
import cartagows.wsframework.wssecurity.policy.model.X509Token;

public class SecurityUtil {

	public static String addWsuIdToElement(OMElement elem) {
		String id;
		OMAttribute idAttr = null;
		//first try to get the Id attr
		idAttr = elem.getAttribute(new QName("Id"));
		if(idAttr == null) {
			//then try the wsu:Id value
			idAttr = elem.getAttribute(new QName(WSConstants.WSU_NS, "Id"));
		}
		if(idAttr == null && elem.getLocalName().equals("Body")){
			idAttr = elem.getFirstElement().getAttribute(new QName("Id"));
			if(idAttr == null) {
				//then try the wsu:Id value
				idAttr = elem.getFirstElement().getAttribute(new QName(WSConstants.WSU_NS, "Id"));
			}
		}
		if(idAttr != null) {
			id = idAttr.getAttributeValue();
		} else {
			//Add an id
			OMNamespace ns = elem.getOMFactory().createOMNamespace(
					WSConstants.WSU_NS, WSConstants.WSU_PREFIX);
			id = "Id-" + elem.hashCode();
			idAttr = elem.getOMFactory().createOMAttribute("Id", ns, id);
			elem.addAttribute(idAttr);
		}

		return id;
	}

	public static Element appendChildToSecHeader(MessageData md,
			Element elem) {
		Element secHeaderElem = md.getSecHeader().getSecurityHeader();
		Node node = secHeaderElem.getOwnerDocument().importNode(
				elem, true);
		return (Element)secHeaderElem.appendChild(node);
	}

	public static Element insertSiblingAfter(MessageData md,
			Element child, Element sibling) {
		if (child == null) {
			return appendChildToSecHeader(md, sibling);
		} else {
			if (child.getOwnerDocument().equals(sibling.getOwnerDocument())) {

				if (child.getParentNode() == null
						&& !child.getLocalName().equals("UsernameToken")) {
					md.getSecHeader().getSecurityHeader().appendChild(child);
				}
				((OMElement) child).insertSiblingAfter((OMElement) sibling);
				return sibling;
			} else {
				Element newSib = (Element) child.getOwnerDocument().importNode(
						sibling, true);
				((OMElement) child).insertSiblingAfter((OMElement) newSib);
				return newSib;
			}
		}
	}

	public static Element insertSiblingBefore(MessageData md, Element child, Element sibling) {
		if(child == null) {
			return appendChildToSecHeader(md, sibling);
		} else {
			if(child.getOwnerDocument().equals(sibling.getOwnerDocument())) {
				((OMElement)child).insertSiblingBefore((OMElement)sibling);
				return sibling;
			} else {
				Element newSib = (Element)child.getOwnerDocument().importNode(sibling, true);
				((OMElement)child).insertSiblingBefore((OMElement)newSib);
				return newSib;
			}
		}

	}

	public static Vector getEncryptedParts(MessageData md) {
		PolicyData pd =  md.getPolicyData();
		SOAPEnvelope envelope = SAAJUtil.getSOAPEnvelopeFromDOOMDocument(md.getDocument());
		Vector encryptedPartsElements  = getPartsAndElements(false, envelope, pd.isEncryptBody(), pd.getEncryptedParts(), pd.getEncryptedElements(),pd.getDeclaredNamespaces());
		return getContentEncryptedElements(encryptedPartsElements, envelope, pd.getContentEncryptedElements(), pd.getDeclaredNamespaces());
	}

	public static Vector getSignedParts(MessageData md) {
		PolicyData pd =  md.getPolicyData();
		SOAPEnvelope envelope = SAAJUtil.getSOAPEnvelopeFromDOOMDocument(md.getDocument());
		return getPartsAndElements(true, envelope, pd.isSignBody(), pd.getSignedParts(), pd.getSignedElements(), pd.getDeclaredNamespaces());
	}

	private static Set findAllPrefixNamespaces(OMElement currentElement, HashMap decNamespacess) {
		Set results = new HashSet();

		//Find declared namespaces
		findPrefixNamespaces(currentElement,results);

		//Get all default namespaces
		List defaultNamespaces = getDefaultPrefixNamespaces(currentElement.getOMFactory());
		for (Iterator iterator = defaultNamespaces.iterator(); iterator
		.hasNext();) {
			OMNamespace ns = (OMNamespace) iterator.next();
			results.add(ns);
		}

		for ( Iterator iterator = decNamespacess.keySet().iterator(); iterator.hasNext();) {
			String prefix  = (String) iterator.next();
			String ns = (String) decNamespacess.get(prefix); 
			OMFactory omFactory = currentElement.getOMFactory();
			OMNamespace namespace = omFactory.createOMNamespace(ns, prefix);
			results.add(namespace);

		}

		return results;
	}

	private static void findPrefixNamespaces(OMElement e, Set results) {

		Iterator iter = e.getAllDeclaredNamespaces();

		if (iter!=null) {
			while (iter.hasNext())
				results.add(iter.next());
		}

		Iterator children = e.getChildElements();

		while (children.hasNext()) {
			findPrefixNamespaces((OMElement)children.next(), results);
		}
	}

	private static List getDefaultPrefixNamespaces(OMFactory factory) {
		List namespaces = new ArrayList();

		// put default namespaces here (sp, soapenv, wsu, etc...)
		namespaces.add(factory.createOMNamespace(WSConstants.ENC_PREFIX, WSConstants.ENC_NS));
		namespaces.add(factory.createOMNamespace(WSConstants.SIG_PREFIX, WSConstants.SIG_NS));
		namespaces.add(factory.createOMNamespace(WSConstants.WSSE_PREFIX, WSConstants.WSSE_NS));
		namespaces.add(factory.createOMNamespace(WSConstants.WSU_PREFIX, WSConstants.WSU_NS));

		return namespaces;

	}

	public static Vector getContentEncryptedElements (Vector encryptedPartsElements, SOAPEnvelope envelope,Vector elements, HashMap decNamespaces ) {

		Set namespaces = findAllPrefixNamespaces(envelope, decNamespaces);

		Iterator elementsIter = elements.iterator();
		while (elementsIter.hasNext())
		{
			String expression = (String)elementsIter.next();
			try {
				XPath xp = new AXIOMXPath(expression);
				Iterator nsIter = namespaces.iterator();

				while (nsIter.hasNext())
				{
					OMNamespace tmpNs = (OMNamespace)nsIter.next();
					xp.addNamespace(tmpNs.getPrefix(), tmpNs.getNamespaceURI());
				}

				List selectedNodes = xp.selectNodes(envelope);

				Iterator nodesIter = selectedNodes.iterator();

				while (nodesIter.hasNext())
				{
					OMElement e = (OMElement)nodesIter.next();

					WSEncryptionPart encryptedElem = new WSEncryptionPart(e.getLocalName(), e.getNamespace().getNamespaceURI(), "Content");
					OMAttribute wsuId = e.getAttribute(new QName(WSConstants.WSU_NS, "Id"));

					if ( wsuId != null ) {
						encryptedElem.setEncId(wsuId.getAttributeValue());
					}

					encryptedPartsElements.add(encryptedElem);

				}

			} catch (JaxenException e) {
				// This has to be changed to propagate an instance of a SecurityException up
				throw new RuntimeException(e);
			}
		}


		return encryptedPartsElements;

	}

	public static Vector getPartsAndElements(boolean sign, SOAPEnvelope envelope, boolean includeBody, Vector parts, Vector elements, HashMap decNamespaces) {

		Vector found = new Vector();
		Vector result = new Vector();

		// check body
		if(includeBody) {
			if( sign ) {
				result.add(new WSEncryptionPart(addWsuIdToElement(envelope.getBody())));
			} else {
				result.add(new WSEncryptionPart(addWsuIdToElement(envelope.getBody()), "Content", WSConstants.PART_TYPE_BODY));
			}
			found.add( envelope.getBody() );
		}

		// Search envelope header for 'parts' from Policy (SignedParts/EncryptedParts)
		SOAPHeader header = envelope.getHeader();

		for(int i=0; i<parts.size(); i++) {
			WSEncryptionPart wsep = (WSEncryptionPart) parts.get( i );
			if( wsep.getName() == null ) {
				// NO name - search by namespace
				ArrayList headerList = header.getHeaderBlocksWithNSURI( wsep.getNamespace() );

				for(int j=0; j<headerList.size(); j++) {
					SOAPHeaderBlock shb = (SOAPHeaderBlock) headerList.get( j ); 

					// find reference in envelope
					OMElement e = header.getFirstChildWithName( shb.getQName() );

					if( ! found.contains(  e ) ) {
						// found new
						found.add( e );

						if( sign ) {
							result.add(new WSEncryptionPart(e.getLocalName(), wsep.getNamespace(), "Content"));
						} else {

							WSEncryptionPart encryptedHeader = new WSEncryptionPart(e.getLocalName(), wsep.getNamespace(), "Content", WSConstants.PART_TYPE_HEADER);
							OMAttribute wsuId = e.getAttribute(new QName(WSConstants.WSU_NS, "Id"));

							if ( wsuId != null ) {
								encryptedHeader.setEncId(wsuId.getAttributeValue());
							}

							result.add(encryptedHeader);
						}
					} 
				}
			} else {
				// try to find
				OMElement e = header.getFirstChildWithName( new QName(wsep.getNamespace(), wsep.getName()) );
				if( e != null ) {
					if( ! found.contains( e ) ) {
						// found new (reuse wsep)
						found.add( e );          
						/*wsep.setType(WSConstants.PART_TYPE_HEADER);
						OMAttribute wsuId = e.getAttribute(new QName(WSConstants.WSU_NS, "Id"));

						if ( wsuId != null ) {
							wsep.setEncId(wsuId.getAttributeValue());
						}

						result.add( wsep );*/
						
						WSEncryptionPart encryptedHeader = new WSEncryptionPart(e.getLocalName(), wsep.getNamespace(), "Content", WSConstants.PART_TYPE_HEADER);
						OMAttribute wsuId = e.getAttribute(new QName(WSConstants.WSU_NS, "Id"));

						if ( wsuId != null ) {
							encryptedHeader.setEncId(wsuId.getAttributeValue());
						}

						result.add(encryptedHeader);
					}
				} 
			} 
		}

		// ?? Search for 'Elements' here

		// decide what exactly is going to be used - only the default namespaces, or the list of all declared namespaces in the message !
		Set namespaces = findAllPrefixNamespaces(envelope, decNamespaces);

		/*Iterator elementsIter = elements.iterator();
		while (elementsIter.hasNext())
		{
			String expression = (String)elementsIter.next();
			try {
				XPath xp = new AXIOMXPath(expression);
				Iterator nsIter = namespaces.iterator();

				while (nsIter.hasNext())
				{
					OMNamespace tmpNs = (OMNamespace)nsIter.next();
					xp.addNamespace(tmpNs.getPrefix(), tmpNs.getNamespaceURI());
				}

				List selectedNodes = xp.selectNodes(envelope);

				Iterator nodesIter = selectedNodes.iterator();
				while (nodesIter.hasNext())
				{
					OMElement e = (OMElement)nodesIter.next();

					if (sign) {
						result.add(new WSEncryptionPart(e.getLocalName(), e.getNamespace().getNamespaceURI(), "Content"));
					} else {

						WSEncryptionPart encryptedElem = new WSEncryptionPart(e.getLocalName(), e.getNamespace().getNamespaceURI(), "Element");
						OMAttribute wsuId = e.getAttribute(new QName(WSConstants.WSU_NS, "Id"));

						if ( wsuId != null ) {
							encryptedElem.setEncId(wsuId.getAttributeValue());
						}

						result.add(encryptedElem);
					}
				}

			} catch (JaxenException e) {
				// This has to be changed to propagate an instance of a SecurityException up
				throw new RuntimeException(e);
			}
		}*/
		for(int i=0;i<elements.size();i++){
			String[] st1 = ((String) elements.get(i)).split("/");
			String[] st2 = st1[st1.length-1].split(":");
			String name = st2[st2.length-1];
			String nm = st2[st2.length-2];
			String namespace = (String) decNamespaces.get(nm);
			WSEncryptionPart encP = new WSEncryptionPart(name, namespace,"Content");
			result.add(encP);
		}

		return result;
	}

	public static boolean checkRequiredElements(SOAPEnvelope envelope, HashMap decNamespaces, String expression ) {

		Set namespaces = findAllPrefixNamespaces(envelope, decNamespaces);

		try {
			XPath xp = new AXIOMXPath(expression);
			Iterator nsIter = namespaces.iterator();

			while (nsIter.hasNext())
			{
				OMNamespace tmpNs = (OMNamespace)nsIter.next();
				xp.addNamespace(tmpNs.getPrefix(), tmpNs.getNamespaceURI());
			}

			List selectedNodes = xp.selectNodes(envelope);

			if (selectedNodes.size() == 0 ) {
				return false;
			}

		} catch (JaxenException e) {
			// This has to be changed to propagate an instance of a SecurityException up
			throw new RuntimeException(e);
		}

		return true;
	}

	public static void setEncryptionUser(MessageData md, WSSecEncryptedKey encrKeyBuilder) throws SecurityException {
		PolicyData pd = md.getPolicyData();
		String encrUser = md.getEncryptionUser();
		if(md.isInitiator()){
			if(encrUser == null || "".equals(encrUser)) {
				throw new SecurityException("missingEncryptionUser");
			} else {
				encrKeyBuilder.setUserInfo(encrUser);
			}
		} else if(!md.isInitiator()){
			Object resultsObj = md.getResults();
			if(resultsObj != null) {
				X509Certificate cert = getReqSigCert((Vector)resultsObj);
				if(cert != null){
					encrKeyBuilder.setUseThisCert(cert);
					//TODO This is a hack, this should not come under USE_REQ_SIG_CERT
					if(encrKeyBuilder.isCertSet()) {
						encrKeyBuilder.setUserInfo(getUsername((Vector)resultsObj));
					}
				} else {
					if(encrUser != null || !"".equals(encrUser)){
						encrKeyBuilder.setUserInfo(encrUser);
					} else {
						throw new SecurityException("missingEncryptionUser");
					}
				}
			}
		}

	}

	public static void setKeyIdentifierType(PolicyData pd, WSSecBase secBase,cartagows.wsframework.wssecurity.policy.model.Token token) {

		if (token.getInclusion() == SPConstants.INCLUDE_TOKEN_NEVER) {

			boolean tokenTypeSet = false;

			if(token instanceof X509Token) {
				X509Token x509Token = (X509Token)token;

				if(x509Token.isRequireIssuerSerialReference()) {
					secBase.setKeyIdentifierType(WSConstants.ISSUER_SERIAL);
					tokenTypeSet = true;
				} else if (x509Token.isRequireKeyIdentifierReference()) {
					secBase.setKeyIdentifierType(WSConstants.SKI_KEY_IDENTIFIER);
					tokenTypeSet = true;
				} else if (x509Token.isRequireThumbprintReference()) {
					secBase.setKeyIdentifierType(WSConstants.THUMBPRINT_IDENTIFIER);
					tokenTypeSet = true;
				}
			} 

			if (!tokenTypeSet) {
				Wss10 wss = pd.getWss11();
				if (wss == null) {
					wss = pd.getWss10();
				}

				if (wss.isMustSupportRefKeyIdentifier()) {
					secBase.setKeyIdentifierType(WSConstants.SKI_KEY_IDENTIFIER);
				} else if (wss.isMustSupportRefIssuerSerial()) {
					secBase.setKeyIdentifierType(WSConstants.ISSUER_SERIAL);
				} else if (wss instanceof Wss11
						&& ((Wss11) wss).isMustSupportRefThumbprint()) {
					secBase.setKeyIdentifierType(WSConstants.THUMBPRINT_IDENTIFIER);
				}
			}

		} else {
			secBase.setKeyIdentifierType(WSConstants.BST_DIRECT_REFERENCE);
		}
	}

	private static X509Certificate getReqSigCert(Vector wsSecEngineResults) {
		/*
		 * Scan the results for a matching actor. Use results only if the
		 * receiving Actor and the sending Actor match.
		 */
		
			for (int j = 0; j < wsSecEngineResults.size(); j++) {
				WSSecurityEngineResult wser =
					(WSSecurityEngineResult) wsSecEngineResults.get(j);
				Integer actInt = (Integer)wser.get(WSSecurityEngineResult.TAG_ACTION);
				if (actInt.intValue() == WSConstants.SIGN) {
					return (X509Certificate)wser.get(WSSecurityEngineResult.TAG_X509_CERTIFICATE);
				}
			}
		

		return null;
	}

	public static String getUsername(Vector results) {
		/*
		 * Scan the results for a matching actor. Use results only if the
		 * receiving Actor and the sending Actor match.
		 */
		for (int i = 0; i < results.size(); i++) {
			WSHandlerResult rResult =
				(WSHandlerResult) results.get(i);

			Vector wsSecEngineResults = rResult.getResults();
			/*
			 * Scan the results for a username token. Use the username
			 * of this token to set the alias for the encryption user
			 */
			for (int j = 0; j < wsSecEngineResults.size(); j++) {
				WSSecurityEngineResult wser =
					(WSSecurityEngineResult) wsSecEngineResults.get(j);
				Integer actInt = (Integer)wser.get(WSSecurityEngineResult.TAG_ACTION);
				if (actInt.intValue() == WSConstants.UT) {
					WSUsernameTokenPrincipal principal = (WSUsernameTokenPrincipal)wser.get(WSSecurityEngineResult.TAG_PRINCIPAL);
					return principal.getName();
				}
			}
		}

		return null;
	}  

	public static String getRequestEncryptedKeyId(Vector results) {

		for (int i = 0; i < results.size(); i++) {
			WSHandlerResult rResult =
				(WSHandlerResult) results.get(i);

			Vector wsSecEngineResults = rResult.getResults();
			/*
			 * Scan the results for the first Signature action. Use the
			 * certificate of this Signature to set the certificate for the
			 * encryption action :-).
			 */
			for (int j = 0; j < wsSecEngineResults.size(); j++) {
				WSSecurityEngineResult wser =
					(WSSecurityEngineResult) wsSecEngineResults.get(j);
				Integer actInt = (Integer)wser.get(WSSecurityEngineResult.TAG_ACTION);
				String encrKeyId = (String)wser.get(WSSecurityEngineResult.TAG_ENCRYPTED_KEY_ID);
				if (actInt.intValue() == WSConstants.ENCR && 
						encrKeyId != null) {
					return encrKeyId;
				}
			}
		}

		return null;
	}

	public static byte[] getRequestEncryptedKeyValue(Vector results) {

		for (int i = 0; i < results.size(); i++) {
			WSHandlerResult rResult =
				(WSHandlerResult) results.get(i);

			Vector wsSecEngineResults = rResult.getResults();
			/*
			 * Scan the results for the first Signature action. Use the
			 * certificate of this Signature to set the certificate for the
			 * encryption action :-).
			 */
			for (int j = 0; j < wsSecEngineResults.size(); j++) {
				WSSecurityEngineResult wser =
					(WSSecurityEngineResult) wsSecEngineResults.get(j);
				Integer actInt = (Integer)wser.get(WSSecurityEngineResult.TAG_ACTION);
				byte[] decryptedKey = (byte[])wser.get(WSSecurityEngineResult.TAG_DECRYPTED_KEY);
				if (actInt.intValue() == WSConstants.ENCR && 
						decryptedKey != null) {
					return decryptedKey;
				}
			}
		}

		return null;
	}

	public static Element insertSiblingAfterOrPrepend(MessageData md, Element child, Element elem) {
		Element retElem = null;
		if(child != null){ // child is not null so insert sibling after
			retElem = SecurityUtil.insertSiblingAfter(md, child, elem);
		}else{ //Prepend 
			retElem = prependSecHeader(md, elem);
		}

		return retElem;
	}

	public static Element insertSiblingBeforeOrPrepend(MessageData md, Element child, Element elem) {
		Element retElem = null;
		if(child != null && child.getPreviousSibling() != null){ 
			retElem = SecurityUtil.insertSiblingBefore(md, child, elem);
		}else{ //Prepend 
			retElem = prependSecHeader(md, elem);
		}

		return retElem;
	}

	private static Element prependSecHeader(MessageData md, Element elem){
		Element retElem = null;

		Element secHeaderElem = md.getSecHeader().getSecurityHeader();
		Node node = secHeaderElem.getOwnerDocument().importNode(
				elem, true);
		Element firstElem = (Element)secHeaderElem.getFirstChild();

		if(firstElem == null){
			retElem = (Element)secHeaderElem.appendChild(node);
		}else{
			if(firstElem.getOwnerDocument().equals(elem.getOwnerDocument())) {
				((OMElement)firstElem).insertSiblingBefore((OMElement)elem);
				retElem = elem;
			} else {
				Element newSib = (Element)firstElem.getOwnerDocument().importNode(elem, true);
				((OMElement)firstElem).insertSiblingBefore((OMElement)newSib);
				retElem = newSib;
			}
		}

		return retElem;
	}

	public static boolean isSecHeaderRequired(PolicyData pd) {

		// Checking for time stamp
		if ( pd.isIncludeTimestamp() ) {
			return true;
		} 

		// Checking for signed parts and elements
		if (pd.isSignBody() || pd.getSignedParts().size() != 0 && 
				pd.getSignedElements().size() != 0) {
			return true;
		}

		// Checking for encrypted parts and elements
		if (pd.isEncryptBody() || pd.getEncryptedParts().size() != 0 && 
				pd.getEncryptedElements().size() != 0 ) {
			return true;
		}   

		// Checking for supporting tokens
		SupportingToken supportingTokens;

		supportingTokens = pd.getSupportingTokens();
		if (supportingTokens != null && supportingTokens.getTokens().size() != 0) {
			return true;
		}

		supportingTokens = pd.getSignedSupportingTokens();
		if (supportingTokens != null && supportingTokens.getTokens().size() != 0) {
			return true;
		}

		supportingTokens = pd.getEndorsingSupportingTokens();
		if (supportingTokens != null && supportingTokens.getTokens().size() != 0) {
			return true;
		}

		supportingTokens = pd.getSignedEndorsingSupportingTokens();
		if (supportingTokens != null && supportingTokens.getTokens().size() != 0) {
			return true;
		}

		return false;

	}

	public static void handleEncryptedSignedHeaders(Vector encryptedParts, Vector signedParts, Document doc) {

		//TODO Is there a more efficient  way to do this ? better search algorithm 
		for (int i = 0 ; i < signedParts.size() ; i++) {
			WSEncryptionPart signedPart = (WSEncryptionPart)signedParts.get(i);

			//This signed part is not a header
			if (signedPart.getNamespace() == null || signedPart.getName() == null) {
				continue;
			}

			for (int j = 0 ; j < encryptedParts.size() ; j ++) {
				WSEncryptionPart encryptedPart = (WSEncryptionPart) encryptedParts.get(j);

				if (encryptedPart.getNamespace() == null || encryptedPart.getName() == null ) {
					continue;
				}

				if (signedPart.getName().equals(encryptedPart.getName()) &&
						signedPart.getNamespace().equals(encryptedPart.getNamespace())) {

					String encDataID =  encryptedPart.getEncId();                    
					Element encDataElem = WSSecurityUtil.findElementById(doc.getDocumentElement(), encDataID, null);

					if (encDataElem != null) {
						Element encHeader = (Element)encDataElem.getParentNode();
						String encHeaderId = encHeader.getAttributeNS(WSConstants.WSU_NS, "Id");

						signedParts.remove(signedPart);
						WSEncryptionPart encHeaderToSign = new WSEncryptionPart(encHeaderId);
						signedParts.add(encHeaderToSign);

					}
				}
			}


		}

	}

	public static String getSigElementId(MessageData md) {

		SOAPEnvelope envelope = SAAJUtil.getSOAPEnvelopeFromDOOMDocument(md.getDocument());

		SOAPHeader header = envelope.getHeader();

		if (header == null ) {
			return null;
		}

		ArrayList secHeaders = header.getHeaderBlocksWithNSURI(WSConstants.WSSE_NS);

		if (secHeaders != null && secHeaders.size() > 0) {
			QName sigQName = new QName(Constants.SignatureSpecNS,Constants._TAG_SIGNATURE);
			QName wsuIdQName = new QName(WSConstants.WSU_NS,"Id");
			OMElement sigElem = ((SOAPHeaderBlock)secHeaders.get(0)).getFirstChildWithName(sigQName);
			OMAttribute wsuId = sigElem.getAttribute(wsuIdQName);

			if (wsuId != null) {
				return wsuId.getAttributeValue();
			}

			wsuId = sigElem.getAttribute(new QName("Id"));

			if (wsuId != null) {
				return wsuId.getAttributeValue();
			}


		}

		return null;
	}

}